"""
Event manager for handling training events in the continual learning framework.

This module provides the core event management infrastructure for the training
pipeline. The EventManager coordinates between event emitters (components that
generate events) and event handlers (components that respond to events) in a
publish-subscribe pattern.

Key Features:
- Priority-based event handler execution
- Event history tracking for debugging and analysis
- Handler enable/disable functionality for dynamic control
- Comprehensive error handling with logging

The event system follows PyTorch Lightning's hook naming conventions and extends
them with continual learning specific events. This allows for clean separation
of concerns and modular handling of training lifecycle events.

Example:
    >>> from learners.events.manager import EventManager
    >>> from learners.events.types import EventType, EventPriority
    >>> from learners.events.context import EventContext
    >>>
    >>> manager = EventManager()
    >>>
    >>> # Register a handler
    >>> def my_handler(context):
    ...     print(f"Handling {context.event.value}")
    ...     return {"handled": True}
    >>>
    >>> manager.register(
    ...     EventType.ON_TRAIN_EPOCH_START,
    ...     my_handler,
    ...     priority=EventPriority.HIGH
    ... )
    >>>
    >>> # Emit an event
    >>> context = EventContext(
    ...     event=EventType.ON_TRAIN_EPOCH_START,
    ...     learner=learner_instance,
    ...     epoch=1
    ... )
    >>> results = manager.emit(context)

Architecture:
    The EventManager maintains separate collections for:
    - Event-specific handlers: Registered for particular event types
    - Event history: Recent events for debugging and analysis

    Handlers are executed in priority order, with error handling ensuring
    that one failing handler doesn't break the entire event flow.
"""

import logging
from collections import defaultdict
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Set, Literal
from typing import overload

from learners.interfaces.event import (
    EventContext,
    EventHandler,
    EventPriority,
    EventType,
    EventManager,
)


@dataclass
class RegisteredHandler:
    """
    Container for registered event handlers with metadata.

    This class wraps event handlers with additional metadata needed
    for proper execution order, identification, and control.

    Attributes:
        handler: The callable event handler function or object
        priority: Execution priority level (lower values execute first)
        name: Human-readable identifier for the handler
        enabled: Whether the handler is currently active
    """

    handler: EventHandler
    priority: EventPriority = EventPriority.NORMAL
    name: Optional[str] = None
    enabled: bool = True


class BaseEventManager(EventManager):
    """
    Central event manager for coordinating training events.

    The EventManager implements a publish-subscribe pattern for handling
    training events throughout the continual learning pipeline. It maintains
    collections of event handlers, executes them in priority order, and
    provides comprehensive event history tracking.

    Key Features:
        - Priority-based handler execution
        - Event history for debugging and analysis
        - Dynamic handler enable/disable
        - Comprehensive error handling

    Attributes:
        _handlers: Event-specific handlers organized by event type
        _event_history: Recent events for debugging (circular buffer)
        _max_history: Maximum number of events to keep in history

    Example:
        >>> manager = EventManager()
        >>>
        >>> # Register handlers
        >>> manager.register(EventType.ON_TASK_START, task_logger)
        >>>
        >>> # Emit events
        >>> context = EventContext(event=EventType.ON_TASK_START, learner=learner)
        >>> results = manager.emit(context)
        >>>
        >>> # Manage handlers
        >>> manager.disable_handler(EventType.ON_TASK_START, "task_logger")
        >>> history = manager.get_event_history(limit=10)
    """

    def __init__(self):
        """
        Initialize the event manager with empty handler collections.

        Sets up the internal data structures for managing event handlers
        and maintaining event history.
        """
        self._handlers: Dict[EventType, List[RegisteredHandler]] = defaultdict(list)
        self._event_history: List[EventContext] = []
        self._max_history = 1000

    def register(
        self,
        event: EventType,
        handler: EventHandler,
        priority: EventPriority = EventPriority.NORMAL,
        name: Optional[str] = None,
    ) -> None:
        """
        Register an event handler for a specific event type.

        Args:
            event: The event type this handler should respond to
            handler: The callable handler (function or object with __call__)
            priority: Execution priority (lower values execute first)
            name: Optional human-readable name for the handler

        Note:
            Handlers are automatically sorted by priority after registration.
            If no name is provided, a unique name is generated from the
            handler's class name and ID.
        """
        registered_handler = RegisteredHandler(
            handler=handler,
            priority=priority,
            name=name or f"{handler.__class__.__name__}_{id(handler)}",
        )

        self._handlers[event].append(registered_handler)
        # Sort by priority
        self._handlers[event].sort(key=lambda x: x.priority.value)

        logging.debug(
            f"Registered handler {registered_handler.name} for event {event.value}"
        )

    def unregister(self, event: EventType, handler: str | EventHandler) -> bool:
        """
        Unregister a specific handler by name.

        Args:
            event: The event type to remove handler from
            handler_name: Name of the handler to remove

        Returns:
            True if handler was found and removed, False otherwise
        """
        handlers = self._handlers.get(event, [])
        handler_name = handler if isinstance(handler, str) else handler.name
        for i, registered_handler in enumerate(handlers):
            if registered_handler.name == handler_name:
                handlers.pop(i)
                logging.debug(
                    f"Unregistered handler {handler_name} for event {event.value}"
                )
                return True
        return False

    def emit(self, context: EventContext) -> Dict[str, Any]:
        """
        Emit an event to all registered handlers for that event type.

        Args:
            context: Event context containing event type and data

        Returns:
            Dictionary mapping handler names to their return values

        Note:
            Handlers are executed in priority order. If a handler raises
            an exception, it is logged but doesn't prevent other handlers
            from executing.
        """
        results = {}

        # Add to history
        self._event_history.append(context)
        if len(self._event_history) > self._max_history:
            self._event_history.pop(0)

        # Execute event-specific handlers
        for registered_handler in self._handlers.get(context.event, []):
            if registered_handler.enabled:
                try:
                    result = registered_handler.handler(context)
                    if result:
                        results[registered_handler.name] = result
                except Exception as e:
                    logging.error(
                        f"Error in handler {registered_handler.name} for event {context.event.value}: {e}"
                    )

        return results

    def enable_handler(self, event: EventType, handler_name: str) -> bool:
        """
        Enable a specific handler by name.

        Args:
            event: The event type the handler is registered for
            handler_name: Name of the handler to enable

        Returns:
            True if handler was found and enabled, False otherwise
        """
        for registered_handler in self._handlers.get(event, []):
            if registered_handler.name == handler_name:
                registered_handler.enabled = True
                return True
        return False

    def disable_handler(self, event: EventType, handler_name: str) -> bool:
        """
        Disable a specific handler by name.

        Args:
            event: The event type the handler is registered for
            handler_name: Name of the handler to disable

        Returns:
            True if handler was found and disabled, False otherwise
        """
        for registered_handler in self._handlers.get(event, []):
            if registered_handler.name == handler_name:
                registered_handler.enabled = False
                return True
        return False

    def get_event_history(self, limit: Optional[int] = None) -> List[EventContext]:
        """
        Get recent event history.

        Args:
            limit: Maximum number of recent events to return.
                  If None, returns all events in history.

        Returns:
            List of EventContext objects in chronological order
        """
        if limit:
            return self._event_history[-limit:]
        return self._event_history.copy()

    def clear_history(self) -> None:
        """Clear event history buffer."""
        self._event_history.clear()

    def clear_handlers(self, event: Optional[EventType] = None) -> None:
        """
        Clear handlers for specific event type or all handlers.

        Args:
            event: Optional specific event type to clear handlers for.
                  If None, clears all handlers.
        """
        if event is None:
            self._handlers.clear()
        else:
            self._handlers[event].clear()

    @overload
    def get_handlers(
        self, event: None = ..., return_dict: Literal[False] = ...
    ) -> Set[EventHandler]: ...

    @overload
    def get_handlers(
        self, event: EventType, return_dict: Literal[False] = ...
    ) -> List[EventHandler]: ...

    @overload
    def get_handlers(
        self, event: EventType | None, return_dict: Literal[True]
    ) -> Dict[str, List[str]]: ...

    def get_handlers(
        self, event: EventType | None = None, return_dict: bool = False
    ) -> Set[EventHandler] | List[EventHandler] | Dict[str, List[str]]:
        """
        Get handlers for a specific event type or all events.

        Args:
            event: Optional specific event type to get handlers for.
                  If None, gets handlers for all event types.
            return_dict: If True, returns handler names instead of handler objects.

        Returns:
            - When return_dict=False and event=None: Set of unique EventHandler objects across all events
            - When return_dict=False and event specified: List of EventHandler objects for that event
            - When return_dict=True and event specified: Dict mapping event name to list of handler names
            - When return_dict=True and event=None: Dict mapping all event names to lists of handler names
        """
        if return_dict:
            if event:
                return {
                    event.value: [
                        h.name
                        for h in self._handlers.get(event, [])
                        if h.name and h.enabled
                    ]
                }

            result = {}
            for evt, handlers in self._handlers.items():
                result[evt.value] = [h.name for h in handlers if h.name and h.enabled]
            return result

        # Return handler objects
        if event is not None:
            return [
                reg_handler.handler
                for reg_handler in self._handlers.get(event, [])
                if reg_handler.enabled
            ]

        return set(
            [
                reg_handler.handler
                for handlers in self._handlers.values()
                for reg_handler in handlers
                if reg_handler.enabled
            ]
        )


# Global event manager instance
event_manager = BaseEventManager()
